# encoding=utf-8
import torch
from thop import profile, clever_format
from ptflops import get_model_complexity_info
from pytorch_model_summary import summary
# load model
from DeeplabModel import DeepLabHead, DeepLabV3, resnet50

# load model
classifier = DeepLabHead(2048, 6)
backbone = resnet50(replace_stride_with_dilation=[False, True, True])
model = DeepLabV3(backbone, classifier, None)
weights_dict = torch.load("model_241.pth", map_location='cpu')['model']
for k in list(weights_dict.keys()):
    if "aux" in k:
        del weights_dict[k]

err_info = model.load_state_dict(weights_dict, strict=False)
print(err_info)
model.eval()

def use_thop():
    test_input = torch.randn(1, 3, 512, 512)
    flops, params = profile(model, inputs=(test_input, ))
    flops, params = clever_format([flops, params], '%.3f')
    print(f"运算量：{flops}, 参数量：{params}")

def use_ptflops():
    input_res = (3, 512, 512)
    macs, params = get_model_complexity_info(model, input_res, as_strings=True, print_per_layer_stat=True)
    print(f"模型 FLOPs: {macs}")
    print(f"模型参数量: {params}")

def use_summary():
    input_tensor = torch.rand(1, 3, 512, 512)
    model_summary = summary(model, input_tensor)
    print(model_summary)

if __name__ == '__main__':
    # use_thop()
    use_ptflops()
    # use_summary()